from src import Preproduction as Pre

pathBase = 'dataset/chat_corpus/clean_chat_corpus/'
paths = [
    "chatterbot.tsv",
    "douban_single_turn.tsv",
    "ptt.tsv",
    "qingyun.tsv",
    "subtitle.tsv",
    "tieba.tsv",
    "weibo.tsv"
]
fullPaths = [pathBase + path for path in paths]

id, ch = Pre.load("results/hanzi_2_one_hot.data")
cnt = len(ch)
V = 10000

print("cnt = " + str(cnt))
print("V = " + str(V))

from src.EmbeddingGloVe import genX as genX
from src.EmbeddingGloVe import Glove as Glove

import numpy as np

X = np.zeros((V, V), dtype=np.float64)
for path in fullPaths:
    print(path)
    X += genX(path, id, V)
np.save("results/X.npy", X)
'''
X = np.load("results/X12000.npy")
X = X[:V, :V]
'''
print(X)

print("Load Finish.\n")

print(X.shape)
print("Non-zeros: ")
print(np.nonzero(X)[0].shape[0])

print("\n")

glove = Glove(V, 128)
#glove.load("results/embedding070.model")

for i in range(1, 101):
    print("Step " + str(i) + ":")
    
    print("Train...")
    loss = glove.trainStep(X)
    print("loss = " + str(loss.data))
    
    c1 = "猫"
    c2 = "幗"
    c3 = "虎"

    print("Eval...")
    i1 = glove.eval(id[c1])
    i2 = glove.eval(id[c2])
    i3 = glove.eval(id[c3])

    d12 = i1 - i2
    d23 = i2 - i3
    d13 = i1 - i3

    print("%s-%s：" % (c1, c2) + str((d12 * d12).sum().data))
    print("%s-%s：" % (c2, c3) + str((d23 * d23).sum().data))
    print("%s-%s：" % (c3, c1) + str((d13 * d13).sum().data))

    if i % 10 == 0:
        glove.save("results/embedding%03d.model", i)

    print("------------------------------")

